#!/usr/bin/env python

import gtk
import sqlite3
import time
import os        
import sys
import random
import cairo
from math import *

class Blob:
    def __init__(self,artist,width,height):
        self.bands = []
        self.artist = artist
        self.width = width
        self.height = height
    
    def __str__(self):
        return self.artist.encode("utf8") +": "+ str(self.bands)
    
    def importance(self):
        position = self.width - self.bands[0][0] # bigger the better
        length = len(self.bands)
        centerdness = 1
        for band in self.bands:
            centerdness += abs((band[1]+band[2]/2)-self.height/2)
        centerdness = self.width*(self.height/2)/centerdness
        #print position/50.0,length*100,centerdness/300.0
        return position/50.0 + length*100 + centerdness/300.0

class Graph:
    WEEK = 7*24*60*60 # seconds
    def __init__(self,user,threshold=10):
        self.user = user
        self.threshold = threshold
        data = self.get_data()
        self.week_totals=data[0]
        self.per_week=data[1]
        """
        weeks = self.per_week.keys()
        weeks.sort()
        for week in weeks:
            print "week",week
            for artist in self.per_week[week].keys():
                print "\t",artist.encode("utf8"),self.per_week[week][artist]
        """
        self.start_week=data[2]
        self.end_week=data[3]
        self.max_week = data[4]
        self.artist_totals = data[5]
        self.handle_cache = {}
    
    def get_data(self):
        db_name = "lastfm-%s.db" % self.user
        self.db_conn = sqlite3.connect(db_name)
        table="artist"
        if not os.path.exists(db_name):
            raise StandardError
        start_week = self._query("SELECT MIN(begin) FROM %ss" % table)[0][0]
        end_week = self._query("SELECT MAX(end) FROM %ss" % table)[0][0]
        totals = {}
        result = self._query("select begin,SUM(playcount) from %ss group by begin" % table) #gets data which is thrown away
        weeks = self._result_to_dict(result).keys()
        weeks.sort()
        per_week = {}
        week_totals = {}
        artist_totals = {}
        max_week = 0
        week = start_week
        week_index = 0
        last_week = start_week - self.WEEK
        two_weeks_ago =  None
        per_week[last_week]={}
        while week < end_week+10:
            per_week[week]={}
            result = self._query("select %s,sum(playcount) from %ss where begin=%s group by %s" % (table,table,week,table)) #potential sql injection
            total=0
            # store playcounts for this week and the artist
            for artist,playcount in result:
                if playcount>=self.threshold:
                    if not artist in artist_totals.keys():
                        artist_totals[artist]=0
                    artist_totals[artist] += playcount
                    per_week[week][artist] = playcount
                    total+=playcount
            # check for artist which started this week
            for artist in per_week[week].keys():
                if not artist in per_week[last_week].keys():
                    per_week[last_week][artist]=0
            # check for artists which ended this week
            for artist in per_week[last_week].keys():
                if not artist in per_week[week].keys() and per_week[last_week][artist]>0:
                    per_week[week][artist]=0
            # check for artist lull last week (added to make graphing easier)
            if two_weeks_ago:
                for artist in per_week[last_week].keys():
                    if (artist in per_week[week].keys() and per_week[week][artist]>0) \
                        and per_week[last_week][artist]==0 \
                        and (artist in per_week[two_weeks_ago].keys() and per_week[two_weeks_ago][artist]>0): #verify lull
                            per_week[last_week][artist]=self.threshold
            week_totals[week] = total
            if total>max_week:
                max_week = total
            two_weeks_ago = last_week
            last_week = week
            if week_index<len(weeks) and weeks[week_index]<=week+self.WEEK+100:
                week = weeks[week_index]
                week_index+=1
            else:
                week += self.WEEK
                #print "manual week increment"
        return (week_totals,per_week,start_week,end_week,max_week,artist_totals)
        
    def draw_decorations(self,context):
        context.set_source_rgb(1, 1, 1)
        context.rectangle(0, 0, self.width, self.height)
        context.fill()
        
        # as logo
        logo = cairo.ImageSurface.create_from_png("as_badge_grey.png")
        context.move_to(0,self.height-50)
        context.set_source_surface(logo)
        context.paint()
        
        #weeks
        for i in range(self.width/self.ppw):
            context.move_to(self.ppw*i,0)
            context.line_to(self.ppw*i,self.height)
        context.set_source_rgba(0, 0, 0,0.1)
        context.set_line_width(1)
        context.stroke()
        
        months = [("Jan",31),("Feb",28),("Mar",31),("Apr",30),("May",31),("Jun",30),("Jul",31),("Aug",31),("Sep",30),("Oct",31),("Nov",30),("Dec",31)]
        #months
        ppd = self.ppw/7.0
        t = time.gmtime(self.start_week)
        month = t[1]%12 # 0 is january for us
        loc = ppd*(months[month][1]-t[2]+11)
        year = t[0]
        context.set_source_rgba(0, 0, 0,0.2)
        context.set_line_width(3)
        context.select_font_face("Sans",cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_NORMAL)
        context.set_font_size(self.ppw/2)
        while loc < self.width:
            context.move_to(loc,0)
            context.line_to(loc,self.height)
            context.stroke()
            
            context.move_to(loc+10,self.height-20)
            context.show_text(unicode(months[month%12][0]+" "+str(year+(month/12))))
            month += 1
            loc += ppd*(months[month%12][1])            
        
    def _get_handle(self,origin,slope,length):
        o_x, o_y = origin
        x = o_x + length/sqrt(1+slope**2)
        y = o_y + length*slope/sqrt(1+slope**2)
        #print "origin",origin,"slope",slope,"length",length,"handle",(x,y)
        return (int(round(x)),int(round(y)))
    
    def _get_cached(self,blob,m_x,forward=True):
        if blob.bands[0][0] in self.current_bottom and self.current_bottom[blob.bands[0][0]]<=blob.bands[0][1]:
            curve = self.current_bottom
        else:
            curve = self.current_top
        weeks = curve.keys()
        weeks.sort()
        if not forward:
            if m_x in weeks and weeks.index(m_x)>0:
                #print "back: cached"
                n_x = weeks[weeks.index(m_x)-1]
                n_y = curve[n_x]
                n_h = 0
            else:
                #print "back: default"
                n_x = m_x - self.ppw
                n_y = self.height/2
                n_h =0
        else:
            if m_x in weeks and weeks.index(m_x)<len(weeks)-1:
                #print "forward: cached"
                n_x = weeks[weeks.index(m_x)+1]
                n_y = curve[n_x]
                n_h = 0
            else:
                #print "forward: default"
                n_x = m_x + self.ppw
                n_y = self.height/2
                n_h =0
        return (n_x,n_y,n_h)
        
    def _get_slope(self,blob,i,bottom=False):
        #return 0
        m_x,m_y,m_h = blob.bands[i]
        if bottom:
            m_y += m_h
        #if (m_x,m_y) in self.slope_cache.keys():
        #    return self.slope_cache[(m_x,m_y)]
        #print blob[i]
        #print last_slope
        if i>=len(blob.bands)-1:# or (bottom and self.current_top.has_key(m_x)):
            n_x,n_y,n_h = self._get_cached(blob,m_x,not bottom)
        else: #get cached point from overall blob
            #print "cached next"
            n_x,n_y,n_h = blob.bands[i+1]
        if i<=0:
            p_x,p_y,p_h = self._get_cached(blob,m_x,bottom)
        else:
            #print "cached previous"
            p_x,p_y,p_h = blob.bands[i-1]
        #print "p:",p_x
        #print "m:",m_x
        #print "n:",n_x
        if bottom:
            p_y += p_h
            n_y += n_h
        slope=(n_y-p_y)/float(n_x-p_x)
        self.slope_cache[(m_x,m_y)]=slope
        return slope
        
    def _get_handles(self,blob,i,bottom=False):
        this_slope = self._get_slope(blob,i,bottom)
        last_slope = self._get_slope(blob,i-1,bottom)
        if bottom:
            sign = -1
        else:
            sign = 1
        p_x,p_y,p_h = blob.bands[i-1]
        m_x,m_y,m_h = blob.bands[i]
        if bottom:
            p_y += p_h
            m_y += m_h
        one = self._get_handle((p_x,p_y),last_slope,sign*self.curve_strength)
        two = self._get_handle((m_x,m_y),this_slope,sign*-1*self.curve_strength)

        last_slope = this_slope
        origin = (m_x,m_y)
        if bottom and (not self.current_bottom.has_key(m_x) or m_y>self.current_bottom[m_x]) and i>1: # this adds to the bottom cache it
            self.current_bottom[m_x] = m_y
        if not bottom and (not self.current_top.has_key(m_x) or m_y<self.current_top[m_x]) and i<len(blob.bands)-1:
            self.current_top[m_x] = m_y
        return (origin,one,two)
    
    def _trace_curve(self,context,blob,bottom=False):
        for i in range(1,len(blob.bands)):
            origin, one, two = self._get_handles(blob,i,bottom)
            context.curve_to(one[0],one[1],two[0],two[1],origin[0],origin[1])
            
    def _draw_handles(self,context,blob,bottom=False):
        for i in range(1,len(blob.bands)):
            p_x = blob.bands[i-1][0]
            p_y = blob.bands[i-1][1]
            p_h = blob.bands[i-1][2]
            if bottom:
                p_y += p_h
            origin, one, two = self._get_handles(blob,i,bottom)
            context.set_source_rgb(0,0,0)
            context.set_line_width(1)
            context.move_to(p_x,p_y)
            context.line_to(one[0],one[1])
            context.stroke()
            context.set_source_rgb(0,0,1)
            context.set_line_width(1)
            context.move_to(origin[0],origin[1])
            context.line_to(two[0],two[1])
            context.stroke()
        
    def draw_blob(self,context,blob,text=None,debug=False):
        #print "blob",blob.artist.encode("utf8")
        #outside part
        if debug:
            self._draw_handles(context,blob)
        context.move_to(blob.bands[0][0],blob.bands[0][1])
        #print (blob[0][0],blob[0][1])
        self._trace_curve(context,blob)
        blob.bands.reverse()
        #print "reverse"
        self._trace_curve(context,blob,True)
        #print "close"
        context.close_path()
        if not self.artist_colors.has_key(blob.artist):
            #self.artist_colors[blob.artist] = (blob.bands[0][1]/float(self.height),1-blob.bands[0][0]/float(self.width),random.random())
            self.artist_colors[blob.artist] = (random.random(),random.random(),random.random())
        r,g,b = self.artist_colors[blob.artist]
        context.set_source_rgba(r, g, b,1)
        context.fill()
        if text:
            context.set_source_rgba(0, 0, 0,1)
            context.select_font_face("Sans",cairo.FONT_SLANT_NORMAL,cairo.FONT_WEIGHT_NORMAL)
            context.set_font_size(10*self.ppp)
            context.move_to(blob.bands[-2][0],blob.bands[-2][1]+blob.bands[-2][2])
            context.show_text(unicode(text))
        if debug:
            self._draw_handles(context,blob,True)

    def get_size(self,ppp=1,ppw=10):
        """ppp is plays per pixel
        ppw is pixels per week"""
        self.width = int(ppw*(4+(self.end_week-self.start_week)/self.WEEK))
        self.height = int(ppp*self.max_week+100)
        self.ppp = ppp
        self.ppw = ppw
        return (self.width,self.height)
    
    def get_blobs(self):
        weeks = self.per_week.keys()
        weeks.sort()
        unfinished_blobs = {} # key is artist, value is a blob (list of x_offset,y_offset,height tuples,top,bottom)
        num_week = 0
        blobs = []
        band=None
        alternator=True
        current_artists=[]
        for week in weeks:
            #print "Processing week:",week
            artists = self.per_week[week].keys()
            artists.sort()
            # add new artists and filter out old
            add_to_top = []
            add_to_bottom = []
            for artist in artists:
                if not artist in current_artists:
                    if alternator:
                        add_to_top.append(artist)
                    else:
                        add_to_bottom.append(artist)
                    alternator=not alternator
            #for a in current_artists lambda a: (a[0],a[1]+1)
            current_artists = add_to_top + filter(lambda a: a in artists,current_artists) + add_to_bottom
            # calculate highest point
            y_offset = self.height/2
            # center index
            """i = 1
            while i<len(current_artists) and (not self.unfinished_blobs.has_key(current_artists[i-1]) or not self.unfinished_blobs.has_key(current_artists[i]) or len(self.unfinished_blobs[current_artists[i]])>=len(self.unfinished_blobs[current_artists[i-1]])):
                y_offset -= self.per_week[week][current_artists[i-1]]*self.ppp
                i+=1
            if i>=len(current_artists):
                y_offset = self.height/2
                if self.week_totals.has_key(week):
                    y_offset -= self.week_totals[week]*self.ppp/2
            else:
                while i<=len(current_artists):
                    y_offset -= self.per_week[week][current_artists[i-1]]*self.ppp/2
                    i+=1"""
            if self.week_totals.has_key(week):
                y_offset -= self.week_totals[week]*self.ppp/2
            last_artist=None
            band=None
            for artist in current_artists:
                top = True
                bottom = True
                # if last band was a start or end
                #if band and band[2]==0:
                #    top = False
                
                # if this band is a start or end
                #if self.per_week[week][artist]==0:
                #    top=False
                #    bottom=False
                    #if band:
                    #    band[4]=False
                
                if not unfinished_blobs.has_key(artist):
                    unfinished_blobs[artist]=Blob(artist,self.width,self.height)
                    blobs.append(unfinished_blobs[artist])
                
                # add the new band
                height = int(round(self.per_week[week][artist]*self.ppp))
                band = [(num_week+1)*self.ppw,y_offset,height]
                unfinished_blobs[artist].bands.append(band)
                y_offset += height

                if len(unfinished_blobs[artist].bands)>1 and self.per_week[week][artist]==0:
                    del unfinished_blobs[artist]
                
            num_week+=1
        return blobs
    
    def blob_draw_order(self,blobs):
        #Separate by week
        sep_by_week = {}
        for blob in blobs:
            if not sep_by_week.has_key(blob.bands[0][0]):
                sep_by_week[blob.bands[0][0]]=[]
            sep_by_week[blob.bands[0][0]].append(blob)
        result = []
        weeks = sep_by_week.keys()
        weeks.sort()
        for week in weeks:
            blobs = sep_by_week[week]
            i = 0
            while i < len(blobs)-1 and blobs[i].bands[0][1]==blobs[i+1].bands[0][1]:
                i+=1
            first_chunk = blobs[:i+1]
            first_chunk.reverse()
            result += first_chunk
            result += blobs[i+1:]
            
        return result
    
    def _draw_current(self,context):
        context.set_source_rgb(1,0,0)
        context.set_line_width(1)
        for x in self.current_top.keys():
            context.move_to(x-5,self.current_top[x])
            context.line_to(x+5,self.current_top[x])
        context.stroke()
        context.set_source_rgb(0,1,0)
        context.set_line_width(1)
        for x in self.current_bottom.keys():
            context.move_to(x-5,self.current_bottom[x])
            context.line_to(x+5,self.current_bottom[x])
        context.stroke()

    def draw(self,context,curve_strength=10):
        self.curve_strength = curve_strength
        if not self.ppp or not self.ppw:
            print "get_size not called first"
            return
        self.draw_decorations(context)
        
        self.artist_colors = {}
        blobs = self.get_blobs()
        blobs = self.blob_draw_order(blobs)
        limit = -1
        start = 12
        i=1
        self.current_top={}
        self.current_bottom={}
        self.slope_cache={}
        for blob in blobs:
            if (start<=i and i<=limit) or limit==-1:
                #print blob
                self.draw_blob(context,blob)#,i,True)#random.random()>0.30)
            i+=1
        #self._draw_current(context)
                
    
    def _query(self,query):
        db_c = self.db_conn.cursor()
        db_c.execute(query)
        result = []
        for row in db_c:
            result.append(row)
        #print result
        return result
    
    def _result_to_dict(self,result):
        d = {}
        for key,value in result:
            d[key]=value
        return d
        
class GtkGraph (gtk.Image):
    def __init__(self,user):
        gtk.Image.__init__(self)
        g = Graph(user,1)
        width, height = g.get_size(ppw=30,ppp=1)
        #print width,height
        colormap = gtk.gdk.colormap_get_system()
        pixmap = gtk.gdk.Pixmap(None,width,height,colormap.get_visual().depth)
        pixmap.set_colormap(colormap)
        self.set_from_pixmap(pixmap,None)
        context = pixmap.cairo_create()
        g.draw(context,10)
        
if __name__=="__main__":
    def destroy(window):
        gtk.main_quit()
    if len(sys.argv) < 2:
        print "Usage: ./lastfm.py <user>"
        sys.exit(1)
    user = sys.argv[1]
    window = gtk.Window()
    window.set_default_size(500,700)
    swindow = gtk.ScrolledWindow()
    widget = GtkGraph(user)
    window.show()
    swindow.show()
    widget.show()
    swindow.add_with_viewport(widget)
    window.add(swindow)
    window.connect("destroy",destroy)
    gtk.main()
